using LinearAlgebra
using StaticArrays
using Printf

# eigen based 
function rotation_vector_eigen(R::SMatrix{3,3,T}) where {T}
    cos_θ      = clamp((tr(R) - 1) / 2, -one(T), one(T))
    vals, vecs = eigen(Matrix(R))
    idx        = argmin(abs.(real.(vals) .- 1))
    n          = real.(vecs[:, idx])
    n          = SVector{3,T}(n[1], n[2], n[3])
    n          = n / norm(n)
    n_skew     = SMatrix{3,3,T}(0, n[3], -n[2], -n[3], 0, n[1], n[2], -n[1], 0)
    sin_θ      = clamp(-tr(n_skew * R) / 2, -one(T), one(T))
    θ          = atan(sin_θ, cos_θ)
    return θ * n
end

# closed form
function rotation_vector_closed(R::SMatrix{3,3,T}) where {T}
    cos_θ = clamp((tr(R) - 1) / 2, -one(T), one(T))

    ax = R[3,2] - R[2,3]
    ay = R[1,3] - R[3,1]
    az = R[2,1] - R[1,2]
    sin_θ = sqrt(ax^2 + ay^2 + az^2) / 2

    if sin_θ > sqrt(eps(T))                       
        θ      = atan(sin_θ, cos_θ)
        inv_2s = θ / (2 * sin_θ)
        return SVector{3,T}(ax * inv_2s, ay * inv_2s, az * inv_2s)

    elseif cos_θ > 0                             
        return zero(SVector{3,T})

    else                                          # θ ≈ π
        if R[1,1] >= R[2,2] && R[1,1] >= R[3,3]
            nx = sqrt(max((R[1,1] + 1) / 2, zero(T)))
            n  = SVector{3,T}(nx, (R[1,2]+R[2,1])/(4nx), (R[1,3]+R[3,1])/(4nx))
        elseif R[2,2] >= R[3,3]
            ny = sqrt(max((R[2,2] + 1) / 2, zero(T)))
            n  = SVector{3,T}((R[1,2]+R[2,1])/(4ny), ny, (R[2,3]+R[3,2])/(4ny))
        else
            nz = sqrt(max((R[3,3] + 1) / 2, zero(T)))
            n  = SVector{3,T}((R[1,3]+R[3,1])/(4nz), (R[2,3]+R[3,2])/(4nz), nz)
        end
        θ = atan(sin_θ, cos_θ)
        return θ * n / norm(n)
    end
end

# rotation vector to rotation matrix using Rodrigues 
function rot_matrix(axis::SVector{3,T}, θ::T) where {T}
    n = axis / norm(axis)
    K = SMatrix{3,3,T}(0, n[3], -n[2], -n[3], 0, n[1], n[2], -n[1], 0)
    return SMatrix{3,3,T}(1I) + sin(θ)*K + (1 - cos(θ))*(K*K)
end

same(a, b; atol) = isapprox(a, b; atol=atol) || isapprox(a, -b; atol=atol)

const T    = Float64
const AXES = [SVector{3,T}(1,0,0), SVector{3,T}(1,1,0), SVector{3,T}(1,-2,3), SVector{3,T}(-1,0.5,2)] # one can choose whatever axis 

ok = true

# Tcheck aggrement for general case
println("CASE 1  eigen vs closed form, θ ∈ [0.1, 3.0]")
@printf("%-16s %-7s %-12s %-12s %-11s\n", "axis", "θ", "‖closed‖", "‖eigen‖", "|Δ|")
worst1 = 0.0
for axn in AXES, θ in (0.1, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0)
    R  = rot_matrix(axn, T(θ)) # def rotation matrix from axis and θ
    vc = rotation_vector_closed(R) # extract rotation vector from closed form method
    ve = rotation_vector_eigen(R) # from eigen method
    d  = min(maximum(abs.(vc .- ve)), maximum(abs.(vc .+ ve)))
    global worst1 = max(worst1, d)
    global ok &= isapprox(norm(vc), θ; atol=1e-8) && same(vc, ve; atol=1e-6)
    @printf("%-16s %-7.2f %-12.6f %-12.6f %-11.2e\n",
            string(round.(axn; digits=1)), θ, norm(vc), norm(ve), d)
end
@printf("\nworst |Δ| (mod sign): %.2e\n\n", worst1)

# Tcheck aggrement for case near θ = 0 
println("CASE 2  near identity")
@printf("%-16s %-7s %-24s %-24s\n", "axis", "θ", "eigen", "closed")
worst2 = 0.0
for axn in AXES, θ in (1e-3, 1e-6, 1e-9)
    R  = rot_matrix(axn, T(θ))
    vc = rotation_vector_closed(R)
    ve = rotation_vector_eigen(R)
    global worst2 = max(worst2, norm(vc - ve), 0.0)
    global ok &= all(isfinite, vc) && all(isfinite, ve)
    @printf("%-16s %-7.0e %-24s %-24s\n",
            string(round.(axn; digits=1)), θ,
            string(round.(ve; digits=10)), string(round.(vc; digits=10)))
end
@printf("\nworst |Δ| near 0: %.2e\n\n", worst2)

# Tcheck aggrement case near θ = π
println("CASE 3  near θ = π : magnitude and axis")
@printf("%-16s %-26s %-26s\n", "true n̂", "eigen n̂", "closed n̂")
for axn in AXES
    n  = axn / norm(axn)
    R  = rot_matrix(axn, T(π) - 1e-7)
    ve = rotation_vector_eigen(R); vc = rotation_vector_closed(R)
    global ok &= isapprox(norm(vc), π; atol=1e-3) && same(vc/norm(vc), SVector{3,T}(n); atol=1e-2)
    @printf("%-16s %-26s %-26s\n",
            string(round.(n; digits=3)),
            string(round.(ve ./ norm(ve); digits=3)),
            string(round.(vc ./ norm(vc); digits=3)))
end
println("\nboth give |Φ| = θ ≈ π and the correct axis.\n")


println(ok ? "PASSED : closed form matches eigen across all cases." :
             "FAILED")